import cv2

tmlpath = 'jinan_tml/'
kmh1 = cv2.imread(tmlpath + 'kmh_1.jpg')
kmh2 = cv2.imread(tmlpath + 'kmh_2.jpg')


def template_demo(tpl, target):
    tpl = cv2.cvtColor(tpl, cv2.COLOR_BGR2GRAY)
    target = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY)
    methods = [cv2.TM_CCOEFF_NORMED]
    th, tw = tpl.shape[:2]
    for md in methods:
        result = cv2.matchTemplate(target, tpl, md)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
        tl = max_loc
        br = (tl[0]+tw, tl[1]+th)
        cv2.rectangle(target, tl, br, (0, 0, 255), 2)
    return max_val, tl, target


def searchno(kmhtel):
    list = []
    pics = ['0.jpg', '1.jpg', '2.jpg', '3.jpg', '4.jpg', '5.jpg', '6.jpg', '7.jpg', '8.jpg', '9.jpg']
    for tmp in pics:
        tmp = cv2.imread(tmlpath + tmp)
        va, lo, fra = template_demo(tmp, kmhtel)
        list.append(va)
    return max(list), str(list.index(max(list))) if max(list) > 0.6 else ''


def getSpeed(img):
    speed = -1
    try:
        im = cv2.resize(img, (352, 288))
        img0 = im[0:30, 170:300]
        val1, loc1, f1 = template_demo(kmh1, img0)
        val2, loc2, f2 = template_demo(kmh2, img0)
        val, loc, f = (val1, loc1, f1) if val1 > val2 else (val2, loc2, f2)
        nmatpic1 = img0[int(loc[1]):int(loc[1] + 12), int(loc[0] - 9):int(loc[0] - 0)]
        nmatpic2 = img0[int(loc[1]):int(loc[1] + 12), int(loc[0] - 18):int(loc[0] - 9)]
        nmatpic3 = img0[int(loc[1]):int(loc[1] + 12), int(loc[0] - 27):int(loc[0] - 18)]

        # ten's digit
        max1, no1 = searchno(nmatpic1)
        # unit's digit
        max2, no2 = searchno(nmatpic2)
        # hundred's digit
        max3, no3 = searchno(nmatpic3)

        speed = int(str(no3 + no2 + no1))
    except Exception as e:
        speed = -999
    return speed


if __name__ == '__main__':
    video = r"../2_.mp4"
    vc = cv2.VideoCapture(video)
    for i in range(1000):
        (ret, frame) = vc.read()
        # 如果没有获取到数据，则结束循环
        if not ret:
            break
        if i % 100 == 0:
            print(getSpeed(frame))

